from cPickle import dump
from cPickle import load
import nltk
import os

# Chunkemoeller

class ChunkStore:
    def __init__(self):
        self.f = None
        self.curdir = os.path.dirname(os.path.abspath(__file__))        
        
    def createFile(self, filename):
        """Creates a file in the current dir, only suplly the filename, not the complete path!"""
        __dir__ = os.path.dirname(os.path.abspath(__file__))
        # filename = "fetch" + str(index) + ".txt"
        filepath = os.path.join(__dir__, filename)
        print filepath
        #f = codecs.open('test', encoding='utf-8', mode='w')
        self.f = open(filepath, 'w')
        return self.f
    
    def saveChunker(self, tagger, chunkername):
        output = open(str(chunkername)+'.pkl', 'wb')
        dump(tagger, output, -1)
        output.close()
        
    def loadChunker(self, chunkername):
        location = self.curdir + "\\" + chunkername
        input = open(str(location)+'.pkl', 'rb')
        tagger = load(input)
        input.close()
        return tagger

        
class ChunkerTrainer:
    def __init__(self):
        pass

    def getUnigramChunker(self):
        chunk_data = [[(t,c) for w,t,c in nltk.chunk.tree2conlltags(chtree)]
                     for chtree in nltk.corpus.conll2000.chunked_sents('train.txt')]
        unigram_chunker = nltk.UnigramTagger(chunk_data)
        print nltk.tag.accuracy(unigram_chunker, chunk_data)
        return unigram_chunker
    

    def getBigramChunker(self):
        chunk_data = [[(t,c) for w,t,c in nltk.chunk.tree2conlltags(chtree)]
                     for chtree in nltk.corpus.conll2000.chunked_sents('train.txt')]
        unigram_chunker = nltk.UnigramTagger(chunk_data)
        bigram_chunker = nltk.BigramTagger(chunk_data, backoff=unigram_chunker)
        print nltk.tag.accuracy(bigram_chunker, chunk_data)
        return bigram_chunker


def flattree(chunks_and_words, full_output=True):
    # [(w, t, c), (w, t, c), ...]
    # full_output = True: [(_w, _org_t, _t, _c)]
    current_chunk = []
    parts = []
    if full_output == False:
        # org tag is not given [(word, tag chunk)]
        words = [w for w, t, c in chunks_and_words]
        tags = [t for w, t, c in chunks_and_words]
        chunks = [c for w, t, c in chunks_and_words]
        org_tags = [None for w, t, c in chunks_and_words]
        chunks_and_words = zip(words, org_tags, tags, chunks)
         
    for w, org_t, t, c in chunks_and_words:
        if full_output:
            # also add the tags to the words
            word = (w, org_t, t)
        else:
            word = w
            
        if c == 'B-NP':
            parts.append(current_chunk)
            current_chunk = ['NP']
            current_chunk.append(word)
        elif c == 'I-NP':
            current_chunk.append(word)
        elif c == 'B-VP':
            parts.append(current_chunk)
            current_chunk = ['VP']            
            current_chunk.append(word)
        elif c == 'B-PP':
            parts.append(current_chunk)
            current_chunk = ['PP']
            current_chunk.append(word)
        elif c == 'O':
            parts.append(current_chunk)
            current_chunk = ['O']
            current_chunk.append(word)
        else:
            current_chunk.append(word)
    parts.append(current_chunk)
    return parts[1:]

def collapseParts(part_list):
    """
        if a PP-chunk is followed by a NP-chunk, the two parts are merged together as a PP-chunk
    
        @param part_list: part_list = [part, part, part, ..]
                          part = [CHUNK, WORD, WORD, ...]
        @return: a part list with collapsed parts
    """
    merged_part_list = []
    pp = None
    for part in part_list:
        if pp != None and part[0] == 'PP':
            merged_part_list.append(pp)
            pp = part
        elif pp == None and part[0] == 'PP':
            pp = part        
        elif pp != None and part[0] == 'NP':
            merged_part = ['PP']
            merged_part.extend(pp[1:])
            merged_part.extend(part[1:])
            merged_part_list.append(merged_part)
            pp = None
        elif pp != None and part[0] != 'NP':
            merged_part_list.append(pp)
            pp = None
            merged_part_list.append(part)
        else: 
            merged_part_list.append(part)
    if pp != None:
        merged_part_list.append(pp)        
    return merged_part_list